# coding=utf-8

# 随机生成layout

# 采用0-1.0
# 随机分栏 1-3
# 先生成 标题 ，数量1-5， 宽度0.5-0.9  高度0.05-0.1
# 公式 2-6 宽度根据分栏情况确定，高度0.02-0.05
# 剩余空间 正文或图片
# 正文最小高度0.2
# 全栏情况下图片最小宽度0.5，多栏情况宽度与栏宽度相同，高度不小于栏宽度的1/2，不大于栏宽度的2倍
# 空隙 0.001-0.01
# 垂直标题 中日韩文，单栏，左右部分插入，高度大于1/3的text旁边， 宽度0.1-0.2

import random
import numpy as np
import cv2
import os
import shutil
# from PIL import Image

vcolumn = random.randint(1,3)
ntitle = random.randint(1,5)
nformula =  0#random.randint(2,6)

W = 1080
H = 1920

language = 'cn'

texthmin = 0.02

def takeSecond(elem):
    return elem[1]

def removeformula(formulalist, titlelist):
    fl = []
    top = 0
    for t in titlelist:
        bot = t[1]
        a = [0,0,0]
        for f in formulalist:
            if f[1]>top and f[1]<bot and a[f[4]]==0:
                a[f[4]] = 1
                fl.append(f)
        top = t[1]+t[3]
    
    bot = 1
    a = [0,0,0]
    for f in formulalist:
        if f[1]>top and f[1]<bot and a[f[4]]==0:
            a[f[4]] = 1
            fl.append(f)
    return fl

def concide(rlist, x, y, w, h):
    if len(rlist)==0:
        return False
    for r in rlist:
        x01 = r[0]
        x02 = r[0]+r[2]
        y01 = r[1]
        y02 = r[1]+r[3]
        x11 = x
        x12 = x+w
        y11 = y
        y12 = y+h
        lx = abs((x01+x02)/2-(x11+x12)/2)
        ly = abs((y01+y02)/2-(y11+y12)/2)
        sax = abs(x01-x02)
        sbx = abs(x11-x12)
        say = abs(y01-y02)
        sby = abs(y11-y12)
        if(lx <= (sax+sbx)/2 and ly <= (say+sby)/2):
            return True
    return False

def splitrect(fr, tir):
    rs = []
    x_t = tir[0]
    w_t = tir[2]

    y_t = tir[1]
    h_t = fr[1]-y_t-random.uniform(0.005, 0.01)
    if(h_t>texthmin):
        print(tir)
        rs.append([x_t, y_t, w_t, h_t, tir[4]])
    
    y_t = fr[1]+fr[3]+random.uniform(0.005, 0.01)
    h_t = tir[1]+tir[3]-y_t
    if(h_t>texthmin):
        rs.append([x_t, y_t, w_t, h_t, tir[4]])
    return rs

def splittextrects(formulalist, tirectlist):
    newtirectlist = []
    for tir in tirectlist:
        needsplit = False
        for fr in formulalist:
            print("DEBUG fr: ", fr)
            if concide([fr], tir[0], tir[1], tir[2], tir[3]):
                newrect = splitrect(fr, tir)
                newtirectlist = newtirectlist+newrect
                needsplit = True
                break
        if not needsplit:
            newtirectlist.append(tir)
    return newtirectlist

def distr(tirectlist):
    titch = [0,0,0,0,0,1,1]#,2]#,3,4] # 2:1:1 text:image:graph:table:tableimage
    textrectlist = []
    imagerectlist = []
    graphrectlist = []
    tablerectlist = []
    tableimgrectlist = []

    for r in tirectlist:
        ratio = r[2]*W/r[3]/H
        tori = int(random.choice(titch))
        if(ratio>0.5 and ratio<2 and tori==1): 
            imagerectlist.append(r)
            print('DBG image#############', r[2], r[3])
        elif(ratio>0.5 and ratio<2 and tori==2): 
            graphrectlist.append(r)
            print('DBG graph#############', r[2], r[3])
        elif(tori==3 and r[3]*H>150 and vcolumn<3 ):
            tablerectlist.append(r)
        elif(tori==4 and r[3]*H>150 and vcolumn<3 ):
            tableimgrectlist.append(r)
        else:
            textrectlist.append(r)

    return textrectlist, imagerectlist, graphrectlist, tablerectlist, tableimgrectlist



img = np.zeros((H,W,3), dtype=np.uint8)
img[:,:,:] = 255

print('colunm: ', vcolumn)

# title

print('ntitle: ', ntitle)
titlelist = []
for i in range(ntitle):
    y_t = random.uniform(0.01, 0.85)
    w_t = random.uniform(0.5, 0.9)
    h_t = random.uniform(0.05, 0.1)
    x_t = random.uniform(0.01, 1-w_t)
    if(not concide(titlelist, x_t, y_t, w_t, h_t)):
        titlelist.append([x_t, y_t, w_t, h_t, 0])

for tr in titlelist:
    x0 = int(tr[0]*W)
    y0 = int(tr[1]*H)
    x1 = int(tr[2]*W)+x0
    y1 = int(tr[3]*H)+y0
    print(x0, y0, x1, y1)
    cv2.rectangle(img, (x0, y0), (x1, y1), (0,0,255), 2)

titlelist.sort(key=takeSecond)

# formula
twoclmn = [0.01, 0.501]
threeclmn = [0.01, 0.34, 0.67]

print('nformula: ', nformula)
formulalist = []
for i in range(nformula):
    y_t = random.uniform(0.01, 0.85)
    clmnidx = 0
    fw = 1.0/vcolumn
    if vcolumn == 2 :
        w_t = fw-0.005
        clmnidx = int(random.choice([0, 1]))
        x_t = twoclmn[clmnidx]
    if vcolumn == 3 :
        w_t = fw-0.005
        clmnidx = int(random.choice([0, 1, 2]))
        x_t = threeclmn[clmnidx]
    if vcolumn == 1 :
        w_t = random.uniform(0.5, 0.8)
        x_t = random.uniform(0.01, 1-w_t)
    h_t = random.uniform(0.02, 0.05)
    if (not concide(titlelist, x_t, y_t, w_t, h_t)) and (not concide(formulalist, x_t, y_t, w_t, h_t)) :
        formulalist.append([x_t, y_t, w_t, h_t, clmnidx])

formulalist.sort(key=takeSecond)
# remove more than 1 formula  # TODO
if vcolumn != 1:
    formulalist = removeformula(formulalist, titlelist)


for tr in formulalist:
    x0 = int(tr[0]*W)
    y0 = int(tr[1]*H)
    x1 = int(tr[2]*W)+x0
    y1 = int(tr[3]*H)+y0
    print(x0, y0, x1, y1)
    cv2.rectangle(img, (x0, y0), (x1, y1), (255,0,0), 2)

knowrectlist = titlelist+formulalist
knowrectlist.sort(key=takeSecond)
for r in knowrectlist:
    print("%.3f, %.3f"%(r[1], r[1]+r[3]))

# text and image regions
tirectlist = []
if vcolumn==1:
    y_t = random.uniform(0.005, 0.02)
    for r in knowrectlist:
        if r[1] - y_t > texthmin:
            x_t = random.uniform(0.005, 0.02)
            w_t = 1-2*x_t
            h_t = r[1] - y_t-random.uniform(0.005, 0.01)
            tirectlist.append([x_t, y_t, w_t, h_t, 0])
        y_t = r[1] + r[3] + random.uniform(0.005, 0.01)
    if 1-y_t > texthmin:
        x_t = random.uniform(0.005, 0.02)
        w_t = 1-2*x_t
        h_t = 1.0 - y_t-random.uniform(0.005, 0.01)
        tirectlist.append([x_t, y_t, w_t, h_t, 0])

if vcolumn==2:
    y_t = random.uniform(0.005, 0.02)
    for r in titlelist:
        if r[1] - y_t > texthmin:
            x_t = random.uniform(0.005, 0.02)
            w_t = 0.5-2*x_t
            h_t = r[1] - y_t-random.uniform(0.005, 0.01)
            tirectlist.append([x_t, y_t, w_t, h_t, 0])
            x_t = 0.5+random.uniform(0.005, 0.02)
            w_t = 0.5-2*random.uniform(0.005, 0.02)
            tirectlist.append([x_t, y_t, w_t, h_t, 1])
        y_t = r[1] + r[3] + random.uniform(0.005, 0.01)
    if 1-y_t > texthmin:
        x_t = random.uniform(0.005, 0.02)
        w_t = 0.5-2*x_t
        h_t = 1.0 - y_t-random.uniform(0.005, 0.01)
        tirectlist.append([x_t, y_t, w_t, h_t, 0])
        x_t = 0.5+random.uniform(0.005, 0.02)
        w_t = 0.5-2*random.uniform(0.005, 0.02)
        tirectlist.append([x_t, y_t, w_t, h_t, 1])

if vcolumn==3:
    y_t = random.uniform(0.005, 0.02)
    for r in titlelist:
        if r[1] - y_t > texthmin:
            x_t = random.uniform(0.005, 0.02)
            w_t = 0.33-2*x_t
            h_t = r[1] - y_t-random.uniform(0.005, 0.01)
            tirectlist.append([x_t, y_t, w_t, h_t, 0])
            x_t = 0.34+random.uniform(0.005, 0.02)
            w_t = 0.33-2*random.uniform(0.005, 0.02)
            tirectlist.append([x_t, y_t, w_t, h_t, 1])
            x_t = 0.68+random.uniform(0.005, 0.02)
            w_t = 0.32-2*random.uniform(0.005, 0.02)
            tirectlist.append([x_t, y_t, w_t, h_t, 2])
        y_t = r[1] + r[3] + random.uniform(0.005, 0.01)
    if 1-y_t > texthmin:
        x_t = random.uniform(0.005, 0.02)
        w_t = 0.33-2*x_t
        h_t = 1.0 - y_t-random.uniform(0.005, 0.01)
        tirectlist.append([x_t, y_t, w_t, h_t, 0])
        x_t = 0.34+random.uniform(0.005, 0.02)
        w_t = 0.33-2*random.uniform(0.005, 0.02)
        tirectlist.append([x_t, y_t, w_t, h_t, 1])
        x_t = 0.68+random.uniform(0.005, 0.02)
        w_t = 0.32-2*random.uniform(0.005, 0.02)
        tirectlist.append([x_t, y_t, w_t, h_t, 2])

if(vcolumn!=1):
    tirectlist = splittextrects(formulalist, tirectlist)

textrectlist, imagerectlist, graphrectlist, tablerectlist, tableimgrectlist = distr(tirectlist)

for tr in textrectlist:
    y0 = int(tr[1]*H)
    y1 = int(tr[3]*H)+y0
    lr = -1
    if(vcolumn<3 and tr[3]>=0.3 and (language=='cn' or language=='jp' or language=='kr') ): # add v title
        if vcolumn==1:
            lr = int(random.choice([0, 1]))
        elif vcolumn ==2:
            if tr[4]==0:
                lr = 0
            else:
                lr = 1

        if lr==0:
            x_vt = random.uniform(0.005, 0.02)
            w_vt = random.uniform(0.1, 0.2)
        else:
            w_vt = random.uniform(0.1, 0.2)
            x_vt = 1-random.uniform(0.005, 0.02)-w_vt
        y_vt = tr[1]
        h_vt = tr[3]
        titlelist.append([x_vt, y_vt, w_vt, h_vt, 0])
        x0 = int(x_vt*W)
        x1 = int(w_vt*W)+x0
        print(x0, y0, x1, y1)
        cv2.rectangle(img, (x0, y0), (x1, y1), (255,0,255), 2)
    
    if(lr==-1):
        x0 = int(tr[0]*W)
        x1 = int(tr[2]*W)+x0
    elif(lr==0):
        x0 = int((x_vt+w_vt+random.uniform(0.005, 0.01))*W)
        x1 = int(tr[2]*W)+int(tr[0]*W)
        tr[0] = x0/W
        tr[2] = (x1-x0)/W
    else:
        x0 = int(tr[0]*W)
        x1 = int((x_vt-random.uniform(0.005, 0.01))*W)
        tr[2] = (x1-x0)/W
    print(x0, y0, x1, y1)
    cv2.rectangle(img, (x0, y0), (x1, y1), (0,255,0), 2)

for tr in imagerectlist:
    x0 = int(tr[0]*W)
    y0 = int(tr[1]*H)
    x1 = int(tr[2]*W)+x0
    y1 = int(tr[3]*H)+y0
    print(x0, y0, x1, y1)
    cv2.rectangle(img, (x0, y0), (x1, y1), (0,255,255), 2)

for tr in graphrectlist:
    x0 = int(tr[0]*W)
    y0 = int(tr[1]*H)
    x1 = int(tr[2]*W)+x0
    y1 = int(tr[3]*H)+y0
    print(x0, y0, x1, y1)
    cv2.rectangle(img, (x0, y0), (x1, y1), (0,255,255), 2)

for tr in tablerectlist:
    x0 = int(tr[0]*W)
    y0 = int(tr[1]*H)
    x1 = int(tr[2]*W)+x0
    y1 = int(tr[3]*H)+y0
    print(x0, y0, x1, y1)
    cv2.rectangle(img, (x0, y0), (x1, y1), (255,255,0), 2)

for tr in tableimgrectlist:
    x0 = int(tr[0]*W)
    y0 = int(tr[1]*H)
    x1 = int(tr[2]*W)+x0
    y1 = int(tr[3]*H)+y0
    print(x0, y0, x1, y1)
    cv2.rectangle(img, (x0, y0), (x1, y1), (255,128,0), 2)

cv2.imwrite('layout.jpg', img)

# save to config
def ConfigGen():
    defaultConfig = './config.default'
    target = './config.py'
    if os.path.exists(target):
        os.remove(target)
    shutil.copy(defaultConfig, target)
    # with open(target, 'a') as f:
    f = open(target, 'a')
    f.write('picWidth = {}\n'.format(W))
    f.write('picHeight = {}\n'.format(H))

    num = len(titlelist)+len(formulalist)+len(textrectlist)+len(imagerectlist)+len(graphrectlist)+len(tablerectlist)+len(tableimgrectlist)
    index = 0

    regionList = []
    for r in titlelist:
        tempDict = {}
        tempDict['attr']= 'title' 
        tempDict['blank'] = 20
        tempDict['lan'] = language
        tempDict['lineSpacing'] = 0
        tempDict['paraSpacing'] = 0
        tempDict['name'] = 'region'+str(index)
        tempDict['x'] = int(r[0]*W)
        tempDict['y'] = int(r[1]*H)
        if(r[0]+r[2]>=1):
            r[2] = 1-r[0]
        if(r[1]+r[3]>=1):
            r[3] = 1-r[1]
        tempDict['h'] = int(r[3]*H)
        tempDict['w'] = int(r[2]*W)
        if(int(r[3]*H)>int(r[2]*W)):
            tempDict['dir'] = 'v'
        else:
            tempDict['dir'] = 'h'
        index = index+1
        regionList.append(tempDict)

    for r in formulalist:
        tempDict = {}
        tempDict['attr']= 'formula' 
        tempDict['name'] = 'region'+str(index)
        tempDict['x'] = int(r[0]*W)
        tempDict['y'] = int(r[1]*H)
        if(r[0]+r[2]>=1):
            r[2] = 1-r[0]
        if(r[1]+r[3]>=1):
            r[3] = 1-r[1]
        tempDict['h'] = int(r[3]*H)
        tempDict['w'] = int(r[2]*W)
        index = index+1
        regionList.append(tempDict)

    for r in textrectlist:
        tempDict = {}
        tempDict['attr']= 'text' 
        tempDict['blank'] = 20
        tempDict['lan'] = language
        tempDict['lineSpacing'] = 0
        tempDict['paraSpacing'] = 0
        tempDict['name'] = 'region'+str(index)
        tempDict['x'] = int(r[0]*W)
        tempDict['y'] = int(r[1]*H)
        if(r[0]+r[2]>=1):
            r[2] = 1-r[0]
        if(r[1]+r[3]>=1):
            r[3] = 1-r[1]
        tempDict['h'] = int(r[3]*H)
        tempDict['w'] = int(r[2]*W)
        tempDict['dir'] = 'h'
        if (language=='cn' or language=='jp' or language=='kr') and int(r[3]*H)>6*50: # todo int(r[3]*H)>6*50 
            tempDict['dir'] = random.choice(['h', 'v'])
        index = index+1
        regionList.append(tempDict)

    for r in imagerectlist:
        tempDict = {}
        tempDict['attr']= 'image' 
        tempDict['name'] = 'region'+str(index)
        tempDict['x'] = int(r[0]*W)
        tempDict['y'] = int(r[1]*H)
        if(r[0]+r[2]>=1):
            r[2] = 1-r[0]
        if(r[1]+r[3]>=1):
            r[3] = 1-r[1]
        tempDict['h'] = int(r[3]*H)
        tempDict['w'] = int(r[2]*W)
        index = index+1
        regionList.append(tempDict)
    
    for r in graphrectlist:
        tempDict = {}
        tempDict['attr']= 'imagegraph' 
        tempDict['name'] = 'region'+str(index)
        tempDict['x'] = int(r[0]*W)
        tempDict['y'] = int(r[1]*H)
        if(r[0]+r[2]>=1):
            r[2] = 1-r[0]
        if(r[1]+r[3]>=1):
            r[3] = 1-r[1]
        tempDict['h'] = int(r[3]*H)
        tempDict['w'] = int(r[2]*W)
        index = index+1
        regionList.append(tempDict)
    
    for r in tablerectlist:
        tempDict = {}
        tempDict['attr']= 'table' 
        tempDict['blank'] = 20
        tempDict['lan'] = language
        tempDict['lineSpacing'] = 0
        tempDict['paraSpacing'] = 0
        tempDict['name'] = 'region'+str(index)
        tempDict['x'] = int(r[0]*W)
        tempDict['y'] = int(r[1]*H)
        if(r[0]+r[2]>=1):
            r[2] = 1-r[0]
        if(r[1]+r[3]>=1):
            r[3] = 1-r[1]
        tempDict['h'] = int(r[3]*H)
        tempDict['w'] = int(r[2]*W)
        index = index+1
        regionList.append(tempDict)

    for r in tableimgrectlist:
        tempDict = {}
        tempDict['attr']= 'tableimg' 
        tempDict['blank'] = 20
        tempDict['lan'] = language
        tempDict['lineSpacing'] = 0
        tempDict['paraSpacing'] = 0
        tempDict['name'] = 'region'+str(index)
        tempDict['x'] = int(r[0]*W)
        tempDict['y'] = int(r[1]*H)
        if(r[0]+r[2]>=1):
            r[2] = 1-r[0]
        if(r[1]+r[3]>=1):
            r[3] = 1-r[1]
        tempDict['h'] = int(r[3]*H)
        tempDict['w'] = int(r[2]*W)
        index = index+1
        regionList.append(tempDict)

    f.write('regionList = {}\n'.format(regionList))
    f.write('numRegion = {}\n'.format(num))
    f.close()
    return regionList, num
    
if __name__ == "__main__":
    ConfigGen()
